import os
import numpy as np
import cv2
from tqdm import tqdm

def create_pupil_iris_mask(
    count_images_dir='./eye_data/count_images',
    noise_mask_dir='./eye_data/noise_masks',
    eyelid_glint_mask_dir='./eye_data/eyelid_glint_masks',
    eyelash_mask_dir='./eye_data/eyelash_masks',
    output_mask_dir='./eye_data/pupil_iris_masks',
    output_img_dir='./eye_data/pupil_iris_images'
):
    """
    批量生成pupil_iris_mask, 并提取该区域内的事件图像
    """
    
    os.makedirs(output_mask_dir, exist_ok=True)
    os.makedirs(output_img_dir, exist_ok=True)

    idx_list = []
    for fname in os.listdir(count_images_dir):
        if fname.startswith('pos_img_') and fname.endswith('.npy'):
            idx = fname[len('pos_img_'):-4]
            idx_list.append(idx)
    idx_list.sort(key=lambda x: int(x))

    for idx in tqdm(idx_list, desc="生成pupil_iris_mask并提取区域事件"):
        try:
            # 读取掩码并确保为二值图像
            noise_mask = 1 - np.load(os.path.join(noise_mask_dir, f'noise_mask_{idx}.npy')).astype(np.uint8)
            eyelid_glint_mask = np.load(os.path.join(eyelid_glint_mask_dir, f'eyelid_glint_mask_{idx}.npy')).astype(np.uint8)
            eyelash_mask = np.load(os.path.join(eyelash_mask_dir, f'eyelash_mask_{idx}.npy')).astype(np.uint8)
            
            # # 确保所有掩码为二值 (0或1)
            # noise_mask = (noise_mask > 0).astype(np.uint8)
            # eyelid_glint_mask = (eyelid_glint_mask > 0).astype(np.uint8)
            # eyelash_mask = (eyelash_mask > 0).astype(np.uint8)
            
            # 合并所有不需要的区域
            all_other_masks = np.clip(noise_mask + eyelid_glint_mask + eyelash_mask, 0, 1)
            
            # 反转得到初始瞳孔虹膜区域
            pupil_iris_mask = 1 - all_other_masks
            
            # # 形态学处理：填充孔洞和去除小区域
            # kernel = np.ones((5, 5), np.uint8)
            # pupil_iris_mask = cv2.morphologyEx(pupil_iris_mask, cv2.MORPH_CLOSE, kernel)
            # pupil_iris_mask = cv2.morphologyEx(pupil_iris_mask, cv2.MORPH_OPEN, kernel)
            
            # # 填充可能的内部孔洞
            # contours, _ = cv2.findContours(pupil_iris_mask, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE)
            # for contour in contours:
            #     cv2.drawContours(pupil_iris_mask, [contour], 0, 1, -1)
            
            # 保存掩码
            np.save(os.path.join(output_mask_dir, f'pupil_iris_mask_{idx}.npy'), pupil_iris_mask)
            cv2.imwrite(os.path.join(output_mask_dir, f'pupil_iris_mask_{idx}.png'), pupil_iris_mask * 255)

            # 读取原始事件图像
            pos_img = np.load(os.path.join(count_images_dir, f'pos_img_{idx}.npy'))
            neg_img = np.load(os.path.join(count_images_dir, f'neg_img_{idx}.npy'))
            all_events_img = pos_img + neg_img

            # 应用掩码，提取pupil_iris区域事件
            pupil_iris_img = all_events_img * pupil_iris_mask
            
            # 保存结果
            np.save(os.path.join(output_img_dir, f'pupil_iris_img_{idx}.npy'), pupil_iris_img)
            
            # 标准化图像以便可视化
            if pupil_iris_img.max() > 0:
                normalized_img = (pupil_iris_img / pupil_iris_img.max() * 255).astype(np.uint8)
            else:
                normalized_img = pupil_iris_img.astype(np.uint8)
            cv2.imwrite(os.path.join(output_img_dir, f'pupil_iris_img_{idx}.png'), normalized_img)
            
        except Exception as e:
            print(f"处理索引 {idx} 时出错: {str(e)}")
            continue